import pandas as pd
from django.core.management.base import BaseCommand
from import_data.models import TrainMedicalData, ValMedicalData

def preprocess_data(data):
    data = data.fillna(0).astype(int)
    data['Дата'] = pd.to_datetime(data['Год'].astype(str) + '-01-01') + pd.to_timedelta(data['Номер недели'].sub(1), unit='W')
    # data.set_index('Дата', inplace=True)
    return data

class Command(BaseCommand):
    help = 'Import data from Excel file'

    def handle(self, *args, **kwargs):
        file_path = '~/data/train.xlsx'
        train_df = pd.read_excel(file_path)
        train_df = preprocess_data(train_df)

        file_path = '~/data/val.xlsx'
        val_df = pd.read_excel(file_path)
        val_df = preprocess_data(val_df)      

        train_df.columns = [
            'year', 'week_number', 'densitometer', 'ct', 'ct_with_cu_1_zone',
            'ct_with_cu_2_or_more_zones', 'mmg', 'mri', 'mri_with_cu_1_zone',
            'mri_with_cu_2_or_more_zones', 'rg', 'fluorograph', 'date'
        ]

        for _, row in train_df.iterrows():
            train_medical_data = TrainMedicalData(
                year=row['year'],
                week_number=row['week_number'],
                densitometer=row['densitometer'],
                ct=row['ct'],
                ct_with_cu_1_zone=row['ct_with_cu_1_zone'],
                ct_with_cu_2_or_more_zones=row['ct_with_cu_2_or_more_zones'],
                mmg=row['mmg'],
                mri=row['mri'],
                mri_with_cu_1_zone=row['mri_with_cu_1_zone'],
                mri_with_cu_2_or_more_zones=row['mri_with_cu_2_or_more_zones'],
                rg=row['rg'],
                fluorograph=row['fluorograph'],
                date=row['date']
            )
            train_medical_data.save()

        val_df.columns = [
            'year', 'week_number', 'densitometer', 'ct', 'ct_with_cu_1_zone',
            'ct_with_cu_2_or_more_zones', 'mmg', 'mri', 'mri_with_cu_1_zone',
            'mri_with_cu_2_or_more_zones', 'rg', 'fluorograph', 'date'
        ]

        for _, row in val_df.iterrows():
            val_medical_data = ValMedicalData(
                year=row['year'],
                week_number=row['week_number'],
                densitometer=row['densitometer'],
                ct=row['ct'],
                ct_with_cu_1_zone=row['ct_with_cu_1_zone'],
                ct_with_cu_2_or_more_zones=row['ct_with_cu_2_or_more_zones'],
                mmg=row['mmg'],
                mri=row['mri'],
                mri_with_cu_1_zone=row['mri_with_cu_1_zone'],
                mri_with_cu_2_or_more_zones=row['mri_with_cu_2_or_more_zones'],
                rg=row['rg'],
                fluorograph=row['fluorograph'],
                date=row['date']
            )
            val_medical_data.save()

        self.stdout.write(self.style.SUCCESS('Data imported successfully'))
